import torch
from ..VideoBenchmark import LVCTestNeural

from tqdm import tqdm

from DCVC.DCVC.src.models.DCVC_net import DCVC_net
from DCVC.DCVC_TCM.src.models.video_net_dmc import DMC as DCVC_TCM
from DCVC.DCVC_DC.src.models.video_model import DMC as DCVC_DC
from DCVC.DCVC_DC.src.utils.stream_helper import get_state_dict as DC_get_state_dict

from DCVC.DCVC_HEM.src.models.video_model import DMC as DCVC_HEM

from EFVC.src.model import EFVC
from EFVC.src.model_distillation import EFVC_Destilation
from EFVC.src.model_OD import EFVC_OD

import numpy as np
import math
from PIL import Image

from util.metrics import compute_msssim, compute_psnr, mse2psnr
from util.io.ImageIO import img2torch
from util.io.VideoIO import RawVideoSequence, YUV420VideoSequence

class DCVCTest(LVCTestNeural):
    CKPT = [
        "../../DCVC/DCVC/checkpoints/model_dcvc_quality_0_psnr.pth",
        "../../DCVC/DCVC/checkpoints/model_dcvc_quality_1_psnr.pth",
        "../../DCVC/DCVC/checkpoints/model_dcvc_quality_2_psnr.pth",
        "../../DCVC/DCVC/checkpoints/model_dcvc_quality_3_psnr.pth"
    ]

    def __init__(self, seqs, ckpt, intra_quality=6, gop=12, **kwargs):
        super().__init__(seqs, intra_quality, gop, **kwargs)

        self.inter_model = DCVC_net()
        self.inter_model.load_state_dict(torch.load(ckpt))
        self.inter_model.eval().cuda()

    def inter_encode(self, input, **kwargs):
        return self.inter_model(self.dpb["__ref"], input)
    
    def inter_eval(self, input, output, **kwargs):
        x_unpadded = self.data_postprocess(input)
        x_hat_unpadded = self.data_postprocess(output["recon_image"])

        psnr = compute_psnr(x_unpadded, x_hat_unpadded)
        msssim = compute_msssim(x_unpadded, x_hat_unpadded)
        bpp = output["bpp"].item() * self.bpp_ratio

        return {
            "x_hat": output["recon_image"],
            "psnr": psnr,
            "msssim": msssim,
            "bpp": bpp,
            "extra": None
        }


class DCVC_TCM_Test(LVCTestNeural):
    def __init__(self, seqs, ckpt, intra_quality = 6, gop = 12, **kwargs):
        super().__init__(seqs, intra_quality, gop, **kwargs)

        self.inter_model = DCVC_TCM()
    
        # state = torch.load(ckpt, map_location='cpu')["state_dict"]
        state = torch.load(ckpt, map_location='cpu')

        self.inter_model.load_state_dict(state)
        self.inter_model.eval().cuda()

        # propagation
        self.feature = None

        self.dpb["feature"] = None


    def inter_encode(self, input, **kwargs):
        output = self.inter_model(input, self.dpb["__ref"], self.dpb["feature"])

        self.dpb["feature"] = output["feature"]

        return output
    
    def inter_eval(self, input, output, **kwargs):
        x_unpadded = self.data_postprocess(input)
        x_hat_unpadded = self.data_postprocess(output["recon_image"])

        psnr = compute_psnr(x_unpadded, x_hat_unpadded)
        msssim = compute_msssim(x_unpadded, x_hat_unpadded)
        bpp = output["bpp"].item() * self.bpp_ratio

        return {
            "x_hat": output["recon_image"],
            "psnr": psnr,
            "msssim": msssim,
            "bpp": bpp,
            "extra": None
        }
    
    
    def intra_encode(self, input_frame: torch.Tensor):
        self.dpb["feature"] = None     # reset feature

        return super().intra_encode(input_frame)


class DCVC_HEM_Test(LVCTestNeural):
    def __init__(self, seqs, ckpt, intra_quality = 6, gop=12, device="cuda", **kwargs):
        super().__init__(seqs, intra_quality, gop, device, **kwargs)

        self.inter_model = DCVC_HEM()
        self.inter_model.load_state_dict(torch.load(ckpt, map_location = "cpu"))
        self.inter_model.eval().cuda()

        print("y_q_scale: ", self.inter_model.y_q_scale.detach().cpu().reshape(4).numpy())
        print("mv_q_scale: ", self.inter_model.mv_y_q_scale.detach().cpu().reshape(4).numpy())

        self.dpb = {
            "ref_frame": None,
            "ref_feature": None,
            "ref_y": None,
            "ref_mv_y": None,
        }

        self.q_index = int(kwargs["q_index"])
    
    def intra_encode(self, input_frame: torch.Tensor):
        out = super().intra_encode(input_frame)

        self.dpb = {
            "ref_frame": out["x_hat"],
            "ref_feature": None,
            "ref_y": None,
            "ref_mv_y": None,
        }

        return out
    
    def inter_encode(self, input, **kwargs):
        result = self.inter_model.forward_one_frame(
            input, self.dpb, 
            mv_y_q_scale = self.inter_model.mv_y_q_scale[self.q_index], 
            y_q_scale = self.inter_model.y_q_scale[self.q_index]
        )

        self.dpb = result["dpb"]

        return result
    
    def inter_eval(self, input, output, **kwargs):
        x_unpadded = self.data_postprocess(input)
        x_hat_unpadded = self.data_postprocess(self.dpb["ref_frame"])

        psnr = compute_psnr(x_unpadded, x_hat_unpadded)
        msssim = compute_msssim(x_unpadded, x_hat_unpadded)
        bpp = output["bpp"].item() * self.bpp_ratio

        return {
            "x_hat": self.dpb["ref_frame"],
            "psnr": psnr,
            "msssim": msssim,
            "bpp": bpp,
            "extra": None
        }
    



class DCVC_DC_Test(LVCTestNeural):
    def __init__(self, seqs, ckpt, intra_quality = 6, gop=12, device="cuda", **kwargs):
        super().__init__(seqs, intra_quality, gop, device)

        y_q_scale_enc, y_q_scale_dec, mv_y_q_scale_enc, mv_y_q_scale_dec = \
            DCVC_DC.get_q_scales_from_ckpt(ckpt)
        print("y_q_scale_enc in inter ckpt: ", end='')
        for q in y_q_scale_enc:
            print(f"{q:.3f}, ", end='')
        print()
        print("y_q_scale_dec in inter ckpt: ", end='')
        for q in y_q_scale_dec:
            print(f"{q:.3f}, ", end='')
        print()
        print("mv_y_q_scale_enc in inter ckpt: ", end='')
        for q in mv_y_q_scale_enc:
            print(f"{q:.3f}, ", end='')
        print()
        print("mv_y_q_scale_dec in inter ckpt: ", end='')
        for q in mv_y_q_scale_dec:
            print(f"{q:.3f}, ", end='')
        print()

        # state_dict = torch.load(ckpt, map_location = "cpu")
        state_dict = DC_get_state_dict(ckpt)
        self.inter_model = DCVC_DC(inplace = True)
        self.inter_model.load_state_dict(state_dict)
        self.inter_model.eval().cuda()

        self.dpb = {
            "ref_frame": None,
            "ref_feature": None,
            "ref_mv_feature": None,
            "ref_y": None,
            "ref_mv_y": None,
        }

        self.q_index = int(kwargs["q_index"])
        self.frame_index = 0

        print("q_index", self.q_index)


    def intra_encode(self, input_frame: torch.Tensor):
        out = super().intra_encode(input_frame)

        self.dpb = {
            "ref_frame": self.data_preprocess(out[0]),
            "ref_feature": None,
            "ref_mv_feature": None,
            "ref_y": None,
            "ref_mv_y": None,
        }

        self.frame_index = 0

        return out
    
    def inter_encode(self, input, **kwargs):
        self.frame_index += 1

        result = self.inter_model(
            input, self.dpb, 
            q_in_ckpt = True,
            q_index = self.q_index,
            frame_idx = self.frame_index % 4
        )
        self.dpb = result["dpb"]

        return result
    
    def inter_eval(self, input, output, **kwargs):
        x_unpadded = self.data_postprocess(input)
        x_hat_unpadded = self.data_postprocess(self.dpb["ref_frame"])

        psnr = compute_psnr(x_unpadded, x_hat_unpadded)
        msssim = compute_msssim(x_unpadded, x_hat_unpadded)
        bpp = output["bpp"].item() * self.bpp_ratio

        return {
            "x_hat": self.dpb["ref_frame"],
            "psnr": psnr,
            "msssim": msssim,
            "bpp": bpp,
            "extra": None
        }


class EFVC_Test(LVCTestNeural):
    def __init__(self, seqs, ckpt, intra_quality = 6, gop = 12, **kwargs):
        super().__init__(seqs, intra_quality, gop)

        self.inter_model = EFVC()
        self.inter_model.load_state_dict(torch.load(ckpt, map_location='cpu'))
        self.inter_model.eval().cuda()

        # propagation
        self.dpb["feature"] = None
    
    def inter_encode(self, input, **kwargs):
        output = self.inter_model(input, self.dpb["__ref"], self.dpb["feature"])

        self.dpb["feature"] = output["feature"]

        return output
    
    def inter_eval(self, input, output, **kwargs):
        x_unpadded = self.data_postprocess(input)
        x_hat_unpadded = self.data_postprocess(output["recon_image"])

        psnr = compute_psnr(x_unpadded, x_hat_unpadded)
        msssim = compute_msssim(x_unpadded, x_hat_unpadded)
        bpp = output["bpp"].item() * self.bpp_ratio

        return {
            "x_hat": output["recon_image"],
            "psnr": psnr,
            "msssim": msssim,
            "bpp": bpp,
            "extra": None
        }


    def intra_encode(self, input_frame: torch.Tensor):
        self.dpb["feature"] = None     # reset feature

        return super().intra_encode(input_frame)

